package study.translation.grammar;

import org.apache.log4j.Logger;
import study.translation.utils.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class GrammarAnalyzer
{
    private static final Logger log = Logger.getLogger( GrammarAnalyzer.class );

    private Grammar grammar;
    private Set<Nonterminal> recursionPath = new LinkedHashSet<Nonterminal>();

    private Map<Rule, Set<List<Terminal>>> ruleFirst = new HashMap<Rule, Set<List<Terminal>>>();
    private Map<Nonterminal, Set<List<Terminal>>> first = new HashMap<Nonterminal, Set<List<Terminal>>>();

    public GrammarAnalyzer( Grammar grammar )
    {
        this.grammar = grammar;
    }

    public void analyze( int k ) throws AnalysisException
    {
        for( Nonterminal n : grammar.getNonterminals() )
        {
            if( !first.containsKey( n ) )
            {
                first.put( n, new HashSet<List<Terminal>>() );
            }
            for( Rule rule : getRules( n ) )
            {
                Set<List<Terminal>> prefixes = expand( rule, k );
                log.debug( "FIRST(" + k + ", " + rule + ") = " + prefixes );
                if( !prefixes.isEmpty() && !first.get( n ).addAll( prefixes ) )
                {
                    throw new AnalysisException( "LL(" + k + ") check failed for " + n );
                }
            }
        }
    }

    public Set<List<Terminal>> getFirst( Nonterminal n )
    {
        return first.get( n );
    }

    public Set<List<Terminal>> getFirst( Rule r )
    {
        return ruleFirst.get( r );
    }

    private Set<List<Terminal>> expand( Rule rule, int k ) throws AnalysisException
    {
        if( !ruleFirst.containsKey( rule ) )
        {
            Set<List<Terminal>> first = new HashSet<List<Terminal>>();
            expand( first, new ArrayList<Terminal>(), rule, k );
            ruleFirst.put( rule, first );
        }
        return ruleFirst.get( rule );
    }

    private void expand( Set<List<Terminal>> result, List<Terminal> prefix, Rule rule, int k ) throws AnalysisException
    {
        List<Symbol> symbols = rule.getRightPart();
        for( int i = 0; i < Math.min( symbols.size(), k ); i++ )
        {
            Symbol s = symbols.get( i );
            if( s instanceof Terminal )
            {
                if( prefix.size() < k )
                {
                    prefix.add( ( Terminal )s );
                    if( prefix.size() == k )
                    {
                        result.add( prefix );
                    }
                }
            }
            else if( s instanceof Nonterminal )
            {
                Nonterminal n = ( Nonterminal )s;
                if( prefix.isEmpty() && recursionPath.contains( n ) )
                {
                    throw new AnalysisException( "Left recursion found: " + StringUtils.join( recursionPath, " -> " ) +
                                                 " -> " + n );
                }
                else if( i == 0 )
                {
                    recursionPath.add( n );
                }
                for( Rule r : getRules( n ) )
                {
                    List<Terminal> copy = new ArrayList<Terminal>( prefix );
                    expand( result, copy, r, k - prefix.size() );
                }
                if( i == 0 )
                {
                    recursionPath.remove( n );
                }
            }
        }
    }

    private Set<Rule> getRules( Nonterminal n ) throws AnalysisException
    {
        Set<Rule> rules = n.getRules();
        if( null == rules )
        {
            throw new AnalysisException( "Could not find rules for " + n );
        }
        return rules;
    }
}
